// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/exec/aggregate.h>
#include <pollux/exec/aggregation_hook.h>
#include <pollux/vector/decoded_vector.h>
#include <pollux/vector/flat_vector.h>
#include <pollux/vector/lazy_vector.h>

namespace kumo::pollux::functions::aggregate {
    template<typename TInput, typename TAccumulator, typename TResult>
    class SimpleNumericAggregate : public exec::Aggregate {
    protected:
        explicit SimpleNumericAggregate(TypePtr resultType) : Aggregate(resultType) {
        }

    public:
        void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
        override {
            extractValues(groups, numGroups, result);
        }

    protected:
        template<typename T>
        static constexpr bool kMayPushdown = !std::is_same_v<T, int128_t> &&
                                             !std::is_same_v<T, Timestamp> && !std::is_same_v<T, UnknownValue>;

        // TData is either TAccumulator or TResult, which in most cases are the same,
        // but for sum(real) can differ.
        template<typename TData = TResult, typename ExtractOneValue>
        void doExtractValues(
            char **groups,
            int32_t numGroups,
            VectorPtr *result,
            ExtractOneValue extractOneValue) {
            POLLUX_CHECK_EQ((*result)->encoding(), VectorEncoding::Simple::FLAT);
            auto vector = (*result)->as<FlatVector<TData> >();
            POLLUX_CHECK(
                vector,
                "Unexpected type of the result vector: {}",
                (*result)->type()->toString());
            POLLUX_CHECK_EQ(vector->elementSize(), sizeof(TData));
            vector->resize(numGroups);
            uint64_t *rawNulls = getRawNulls(vector);
            if constexpr (std::is_same_v<TData, bool>) {
                uint64_t *rawValues = vector->template mutableRawValues<uint64_t>();
                for (int32_t i = 0; i < numGroups; ++i) {
                    char *group = groups[i];
                    if (isNull(group)) {
                        vector->set_null(i, true);
                    } else {
                        clearNull(rawNulls, i);
                        bits::setBit(rawValues, i, extractOneValue(group));
                    }
                }
            } else {
                TData *rawValues = vector->mutableRawValues();
                for (int32_t i = 0; i < numGroups; ++i) {
                    char *group = groups[i];
                    if (isNull(group)) {
                        vector->set_null(i, true);
                    } else {
                        clearNull(rawNulls, i);
                        rawValues[i] = extractOneValue(group);
                    }
                }
            }
        }

        // TData is used to store the updated group states. It can be either
        // TAccumulator or TResult, which in most cases are the same, but for
        // sum(real) can differ. TValue is used to decode the update input 'args'.
        // It can be either TAccumulator or TInput, which is most cases are the same
        // but for sum(real) can differ.
        template<
            bool tableHasNulls,
            typename TData = TResult,
            typename TValue = TInput,
            typename UpdateSingleValue>
        void updateGroups(
            char **groups,
            const SelectivityVector &rows,
            const VectorPtr &arg,
            UpdateSingleValue updateSingleValue,
            bool mayPushdown) {
            DecodedVector decoded(*arg, rows, !mayPushdown);
            auto encoding = decoded.base()->encoding();
            if constexpr (kMayPushdown<TData>) {
                if (encoding == VectorEncoding::Simple::LAZY &&
                    !arg->type()->isDecimal()) {
                    pollux::aggregate::SimpleCallableHook<TData, UpdateSingleValue> hook(
                        exec::Aggregate::offset_,
                        exec::Aggregate::nullByte_,
                        exec::Aggregate::nullMask_,
                        groups,
                        &this->exec::Aggregate::numNulls_,
                        updateSingleValue);

                    auto indices = decoded.indices();
                    decoded.base()->as<const LazyVector>()->load(
                        RowSet(indices, arg->size()), &hook);
                    return;
                }
            }

            if (decoded.isConstantMapping()) {
                if (!decoded.is_null_at(0)) {
                    auto value = decoded.value_at<TValue>(0);
                    rows.applyToSelected([&](vector_size_t i) {
                        updateNonNullValue<tableHasNulls, TData>(
                            groups[i], TData(value), updateSingleValue);
                    });
                }
            } else if (decoded.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (decoded.is_null_at(i)) {
                        return;
                    }
                    updateNonNullValue<tableHasNulls, TData>(
                        groups[i], TData(decoded.value_at<TValue>(i)), updateSingleValue);
                });
            } else if (decoded.isIdentityMapping() && !std::is_same_v<TValue, bool>) {
                auto data = decoded.data<TValue>();
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue<tableHasNulls, TData>(
                        groups[i], TData(data[i]), updateSingleValue);
                });
            } else {
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue<tableHasNulls, TData>(
                        groups[i], TData(decoded.value_at<TValue>(i)), updateSingleValue);
                });
            }
        }

        // TData is used to store the updated group state. It can be either
        // TAccumulator or TResult, which in most cases are the same, but for
        // sum(real) can differ. TValue is used to decode the update input 'args'.
        // It can be either TAccumulator or TInput, which is most cases are the same
        // but for sum(real) can differ.
        template<
            typename TData = TResult,
            typename TValue = TInput,
            typename UpdateSingle,
            typename UpdateDuplicate>
        void updateOneGroup(
            char *group,
            const SelectivityVector &rows,
            const VectorPtr &arg,
            UpdateSingle updateSingleValue,
            UpdateDuplicate updateDuplicateValues,
            bool /*mayPushdown*/,
            TData initialValue) {
            DecodedVector decoded(*arg, rows);

            // Do row by row if not all rows are selected.
            if (decoded.isConstantMapping()) {
                if (!decoded.is_null_at(0)) {
                    updateDuplicateValues(
                        initialValue,
                        TData(decoded.value_at<TValue>(0)),
                        rows.countSelected());
                    updateNonNullValue<true, TData>(group, initialValue, updateSingleValue);
                }
            } else if (decoded.may_have_nulls()) {
                rows.applyToSelected([&](vector_size_t i) {
                    if (decoded.is_null_at(i)) {
                        return;
                    }
                    updateNonNullValue<true, TData>(
                        group, TData(decoded.value_at<TValue>(i)), updateSingleValue);
                });
            } else if (decoded.isIdentityMapping() && !std::is_same_v<TValue, bool>) {
                auto data = decoded.data<TValue>();
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue<true, TData>(
                        group, TData(data[i]), updateSingleValue);
                });
            } else {
                rows.applyToSelected([&](vector_size_t i) {
                    updateNonNullValue<true, TData>(
                        group, TData(decoded.value_at<TValue>(i)), updateSingleValue);
                });
            }
        }

        template<typename THook>
        void
        pushdown(char **groups, const SelectivityVector &rows, const VectorPtr &arg) {
            DecodedVector decoded(*arg, rows, false);
            const vector_size_t *indices = decoded.indices();
            THook hook(
                exec::Aggregate::offset_,
                exec::Aggregate::nullByte_,
                exec::Aggregate::nullMask_,
                groups,
                &this->exec::Aggregate::numNulls_);
            // The decoded vector does not really keep the info from the 'rows', except
            // for the 'upper bound' of it. In case not all rows are selected we need to
            // generate proper indices, which we 'indirect' through the ones we got from
            // the decoded vector.
            vector_size_t numIndices{arg->size()};
            if (not rows.isAllSelected()) {
                const auto numSelected = rows.countSelected();
                if (numSelected != arg->size()) {
                    pushdownCustomIndices_.resize(numSelected);
                    vector_size_t tgtIndex{0};
                    rows.applyToSelected([&](vector_size_t i) {
                        pushdownCustomIndices_[tgtIndex++] = indices[i];
                    });
                    indices = pushdownCustomIndices_.data();
                    numIndices = numSelected;
                }
            }

            decoded.base()->as<const LazyVector>()->load(
                RowSet(indices, numIndices), &hook);
        }

    private:
        // TData is either TAccumulator or TResult, which in most cases are the same,
        // but for sum(real) can differ.
        template<
            bool tableHasNulls,
            typename TDataType = TAccumulator,
            typename Update>
        inline void
        updateNonNullValue(char *group, TDataType value, Update updateValue) {
            if constexpr (tableHasNulls) {
                exec::Aggregate::clearNull(group);
            }
            updateValue(*exec::Aggregate::value<TDataType>(group), value);
        }
    };
} // namespace kumo::pollux::functions::aggregate
